# -*- coding: utf-8 -*-

import os
import uuid
from flask import Flask, request, render_template_string, jsonify, Response
import dashscope
from dashscope import TextEmbedding, Generation
from dashvector import Client

app = Flask(__name__)
app.secret_key = uuid.uuid4().hex

# HTML 模板
HTML_TEMPLATE = '''
<!DOCTYPE html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>基于专属知识的问答服务</title>
    <style>
      * {
        margin: 0;
        box-sizing: border-box;
      }
      body {
        font-family: Arial, sans-serif;
        color: #333333;
      }
      .container {
        margin: 20px auto;
        padding: 20px;
        width: 600px;
        height: calc(100vh - 40px);
        display: flex;
        flex-direction: column;
        align-items: center;
        background-color: #f7f8fc;
        border-radius: 20px;
        box-shadow: 0px 0px 30px #cccccc;
      }
      .title {
        align-self: flex-start;
      }
      .content {
        width: 100%;
        margin: 24px 0;
        flex: 1;
        overflow-y: auto;
      }
      .chatItem {
        margin-bottom: 20px;
        display: flex;
      }
      .chatItem:last-child {
        margin-bottom: 0;
      }
      .avatar {
        width: 28px;
        height: 28px;
        border-radius: 50%;
        display: flex;
        justify-content: center;
        align-items: center;
        flex-shrink: 0;
      }
      .self .avatar {
        background: linear-gradient(to bottom, #995ced, #487cfd);
      }
      .customer .avatar {
        background: #ffffff;
      }
      .text {
        margin-left: 16px;
        border-radius: 12px;
        line-height: 24px;
        max-width: 100%;
        padding: 8px 12px;
        white-space: pre-wrap;
      }
      .text img {
        height: 6px;
        object-fit: contain;
      }
      .self .text {
        background-color: #dae7f9;
      }
      .customer .text {
        background-color: #ffffff;
      }
      .questionForm {
        padding: 16px 8px 16px 20px;
        border: 1px solid #e8eaf2;
        border-radius: 12px;
        box-shadow: 0 12px 24px -16px rgba(54, 54, 73, 0.04), 0 12px 40px 0 rgba(51, 51, 71, 0.08), 0 0 1px 0 rgba(44, 44, 54, 0.02);
        width: 100%;
        display: flex;
        align-items: center;
        background: #fff;
      }
      .questionForm:hover {
        border: 1px solid #615ced;
      }
      .question {
        flex: 1;
        border: none;
        appearance: none;
        user-select: auto;
        resize: none;
        outline: none;
        font-size: 16px;
        color: #3f3f3f;
      }
      .question::placeholder {
        color: rgba(135, 138, 171, 0.8);
      }

      .submit {
        width: 36px;
        height: 36px;
        border-radius: 18px;
        background: linear-gradient(0deg, #615ced -2%, #534ad1 93%);
        display: flex;
        justify-content: center;
        align-items: center;
        color: #fff;
        cursor: pointer;
      }
      .disabled {
        cursor: no-drop;
        background: #d6d5de;
      }
    </style>
    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
  </head>
  <body>
    <div class="container">
      <div class="content" id="content"></div>
      <div class="questionForm">
        <textarea placeholder="请将您遇到的问题告诉我" id="question" class="question"></textarea>
        <div id="submit" class="submit disabled">
          <svg t="1729757427561" class="icon" viewBox="0 0 1024 1024" version="1.1" xmlns="http://www.w3.org/2000/svg" p-id="7934" xmlns:xlink="http://www.w3.org/1999/xlink" width="20" height="20">
            <path
              d="M955 125.6c-4.5-12.5-15.4-21.8-28.6-24.2-9.2-1.7-18.4 0.1-26.1 4.8L83.9 544.3c-12.8 6.8-20.6 20.6-19.8 35.1 0.8 14.3 9.9 27.3 23.2 32.9l238.5 97.9c12.4 5.2 26.8 3.4 37.5-4.9s16.1-21.3 14.4-34.8c-1.6-13.2-10.4-24.6-23-29.9l-165-67.7 573-307.3-357 421.9c-6.5 7.7-9.6 17.4-8.8 27.4L411.4 884c1.6 19.7 17.7 34.6 37.5 34.6 9.5 0 18.7-3.7 25.8-10.4l74-69.2 0.1-0.1c13.5-13.2 15.2-33.5 4.8-48.4l222.6 72c4 1.3 7.9 2 11.7 2 18.2 0 33.8-12.9 37-30.6l131.6-688.3h-0.1c1.4-6.6 0.9-13.5-1.4-20zM497.3 783.8L479.9 800l-6.5-75.9L856 271.6l-96.8 506.3L539 706.6c-19.7-6.4-41.1 4.4-47.5 24.2-5.8 17.9 2.5 37.1 18.9 45.4-4.7 1.5-9.2 4-13.1 7.6z"
              p-id="7935"
              fill="currentColor"
              fillOpacity="1"
            ></path>
          </svg>
        </div>
      </div>
    </div>
    <script>
      var host = ""; // 此处修改服务端地址

      var submitBtn = document.getElementById("submit");
      var questionDom = document.getElementById("question");

      function createSelfElement(text) {
        return `<div class="chatItem self">
          <div class="avatar">
            <img src="https://img.alicdn.com/imgextra/i2/O1CN01e3noUR1L6Y5NO3AMa_!!6000000001250-55-tps-20-20.svg" width="14" height="14"/>
          </div>
          <div class="text">${text}</div>
        </div>`;
      }

      function createCustomerElement() {
        return `<div class="chatItem customer">
          <div class="avatar">
            <img src="https://img.alicdn.com/imgextra/i2/O1CN01ZVWmQA25hcT3uGfH3_!!6000000007558-55-tps-28-28.svg" />
          </div>
          <div class="text customerText"><img src="https://img.alicdn.com/imgextra/i4/O1CN01GVqCax1TKIuGGysPl_!!6000000002363-54-tps-416-96.apng" /></div>
        </div>`;
      }

      function scrollToBottom(element) {
        element.scrollTop = element.scrollHeight - element.clientHeight;
      }

      function submitMessage() {
        if (submitBtn.classList.contains("disabled")) return;
        var contentDom = document.getElementById("content");
        var question = questionDom.value;
        var selfStr = createSelfElement(question);
        contentDom.innerHTML += selfStr;
        questionDom.value = "";
        questionDom.setAttribute("disabled", "disabled");

        var eventSource = new EventSource(`${host}/events?question=${encodeURIComponent(question)}`);

        var customerStr = createCustomerElement();
        contentDom.innerHTML += customerStr;

        var customerTexts = document.getElementsByClassName("customerText");
        var lastCustomerText = customerTexts[customerTexts.length - 1];

        eventSource.addEventListener("success", function (event) {
          console.log("event", event.data);
          lastCustomerText.innerHTML = marked.parse(event.data);
          scrollToBottom(contentDom);
        });

        eventSource.addEventListener("stop", function (event) {
          console.log("event", event);
          lastCustomerText.innerHTML = marked.parse(event.data);
          scrollToBottom(contentDom);
          eventSource.close();
          questionDom.removeAttribute("disabled");
        });

        eventSource.onerror = function (err) {
          eventSource.close();
          questionDom.removeAttribute("disabled");
        };
      }

      questionDom.oninput = function () {
        if (this.value) {
          submitBtn.classList.remove("disabled");
        } else {
          submitBtn.classList.add("disabled");
        }
      };

      submitBtn.onclick = submitMessage;
      questionDom.onkeydown = function (event) {
        if (event.key === "Enter") {
          event.preventDefault();
          submitMessage();
        }
      };
    </script>
  </body>
</html>
'''


@app.route('/', methods=['GET', 'POST'])
def home():
    return render_template_string(HTML_TEMPLATE)


@app.route('/events', methods=['GET'])
def events():
    question = request.args.get('question')
    if not question:
        return jsonify({'error': 'Missing required parameter: question'}), 400

    context = get_answers(question)
    return answer_question(question, context)


def get_answers(question):
    dashscope.api_key = os.environ.get('DASHSCOPE_API_KEY')
    # question = '海南安定追尾事故，发生在哪里？原因是什么？人员伤亡情况如何？'
    return search_relevant_news(question)


def search_relevant_news(question):
    # 初始化 dashvector client
    client = Client(
        api_key=os.environ.get('DASHVECTOR_API_KEY'),
        endpoint=os.environ.get('DASHVECTOR_CLUSTER_ENDPOINT')
    )

    # 获取刚刚存入的集合
    collection = client.get('news_embeddings')
    assert collection, collection

    # 向量检索：指定 topk = 1
    rsp = collection.query(generate_embeddings(question), output_fields=['raw'], topk=1)
    assert rsp, rsp
    return rsp.output[0].fields['raw']


def generate_embeddings(news):
    rsp = TextEmbedding.call(
        model=TextEmbedding.Models.text_embedding_v1,
        input=news
    )
    embeddings = [record['embedding'] for record in rsp.output['embeddings']]
    return embeddings if isinstance(news, list) else embeddings[0]


def answer_question(question, context):
    def generate():
        prompt = f'''请基于```以下内容回答问题。
        ```
        {context}
        ```
        我的问题是：{question}。
        '''
        i = 0  # 初始化计数器
        for rsp in Generation.call(model='qwen-turbo', prompt=prompt, stream=True):
            text = rsp['output']['text'].replace("\n", "<br/>")
            finish_reason = rsp['output']['finish_reason'] != 'null'
            message = f"id: {i + 1}\nevent: {finish_reason and 'stop' or 'success'}\ndata: {text}\n\n"
            print(message)
            yield message
            i += 1

    return Response(generate(), mimetype='text/event-stream')


if __name__ == '__main__':
    app.run(host='0.0.0.0', debug=False)
